#James M. Coughlan
#new version: incorporates speed-up trick
#note: the SS version works on messages padded on each side to eliminate special border cases

import time

def update_msgs(msgsX, nd, rad, unPots, binPots, f0, i0, j0, I0, J0, msgs1, I1, J1, msgs2, I2, J2):
    unp = unPots[i0-1][j0-1]
    mX = msgsX[i0+I0][j0+J0]
    m1 = msgs1[i0+I1][j0+J1]
    m2 = msgs2[i0+I2][j0+J2]
    sum_h = 0.
    for d in range(nd):
        sum_h += unp[d]*mX[d]*m1[d]*m2[d]
    msgs = msgsX[i0][j0]
    A=f0*sum_h
    somme = 0.
    for d1 in range(nd):
        new = A
        for d in range(max(0,d1-rad),min(d1+rad+1,nd)):
            new += (binPots[d][d1]-f0)*(unp[d]*mX[d]*m1[d]*m2[d])
        msgs[d1] = new
        somme += new
    for d in range(nd):
        msgs[d] /= somme

def do_sweepsSS2(unPots, binPots, msgsL, msgsR, msgsU, msgsD, nsweeps, h2, w2, nd, rad, f0): # shedskin version
    t0 = time.time()
    for sweep in range(nsweeps):
        print('sweep number:', sweep)

        # do Left update:
        for i0 in range(1,h2-1):
            for j0 in range(w2-2,1,-1):
                update_msgs(msgsL, nd, rad, unPots, binPots, f0, i0, j0, 0, 1, msgsU, 1, 0, msgsD, -1, 0)

        # do Right update:
        for i0 in range(1,h2-1):
            for j0 in range(1,w2-2):
                update_msgs(msgsR, nd, rad, unPots, binPots, f0, i0, j0, 0, -1, msgsU, 1, 0, msgsD, -1, 0)

        # do Up update:
        for j0 in range(1,w2-1):
            for i0 in range(h2-2,1,-1):
                update_msgs(msgsU, nd, rad, unPots, binPots, f0, i0, j0, 1, 0, msgsL, 0, 1, msgsR, 0, -1)

        # do Down update:
        for j0 in range(1,w2-1):
            for i0 in range(1,h2-2):
                update_msgs(msgsD, nd, rad, unPots, binPots, f0, i0, j0, -1, 0, msgsL, 0, 1, msgsR, 0, -1)

    print('sweep in %.2f' % (time.time()-t0))
    return msgsL, msgsR, msgsU, msgsD

if __name__ == '__main__':
    h2,w2,nd,rad,f0=5,6,4,1,0.2
    nsweeps=10
    unPots=[[[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 1.        ,  1.        ,  1.        ,  1.        ]],

    [[ 1.        ,  0.04978707,  0.04978707,  0.04978707],
    [ 0.24659696,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.24659696,  0.04978707,  0.04978707,  0.04978707],
    [ 0.44932896,  0.44932896,  0.44932896,  0.44932896]],

    [[ 1.        ,  0.67032005,  0.67032005,  0.67032005],
    [ 0.06081006,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.81873075,  0.44932896,  0.44932896,  0.44932896],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707]],

    [[ 0.36787944,  0.04978707,  0.04978707,  0.04978707],
    [ 0.81873075,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707]],

    [[ 1.        ,  0.04978707,  0.04978707,  0.04978707],
    [ 0.11080316,  0.04978707,  0.04978707,  0.04978707],
    [ 0.54881164,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707],
    [ 0.81873075,  0.04978707,  0.04978707,  0.04978707],
    [ 0.04978707,  0.04978707,  0.04978707,  0.04978707]]]

    binPots=[[ 1.        ,  0.60653066,  0.36787944,  0.22313016],
    [ 0.60653066,  1.        ,  0.60653066,  0.36787944],
    [ 0.36787944,  0.60653066,  1.        ,  0.60653066],
    [ 0.22313016,  0.36787944,  0.60653066,  1.        ]]

    msgsL=[[[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.29002292,  0.2989277 ,  0.24252004,  0.16852933],
    [ 0.26310106,  0.29748655,  0.25634478,  0.18306761],
    [ 0.26094932,  0.29628075,  0.25764968,  0.18512025],
    [ 0.2660311 ,  0.29309694,  0.2554316 ,  0.18544036],
    [ 0.2429187 ,  0.28254909,  0.26607192,  0.20846029]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.37767773,  0.28994844,  0.20252485,  0.12984898],
    [ 0.27528358,  0.30773349,  0.24927874,  0.16770419],
    [ 0.28793716,  0.30515927,  0.24350915,  0.16339442],
    [ 0.34116812,  0.28865721,  0.22147385,  0.14870082],
    [ 0.24021786,  0.28857552,  0.26886635,  0.20234026]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.35925421,  0.30137138,  0.20782554,  0.13154887],
    [ 0.30184483,  0.30957441,  0.23491658,  0.15366418],
    [ 0.31329719,  0.30255431,  0.23159436,  0.15255414],
    [ 0.28215613,  0.30507813,  0.24650166,  0.16626408],
    [ 0.2502795 ,  0.29243356,  0.26351259,  0.19377434]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.43952647,  0.2797989 ,  0.17408407,  0.10659055],
    [ 0.31874805,  0.30821492,  0.22649261,  0.14654442],
    [ 0.29846323,  0.30843893,  0.23701008,  0.15608777],
    [ 0.2975981 ,  0.30546019,  0.23828072,  0.15866099],
    [ 0.2534437 ,  0.29230677,  0.26173494,  0.19251459]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.38851916,  0.28944442,  0.19638156,  0.12565486],
    [ 0.41445041,  0.28153079,  0.1859509 ,  0.1180679 ],
    [ 0.30207854,  0.30005682,  0.23625284,  0.1616118 ],
    [ 0.40847836,  0.27925759,  0.18936394,  0.1229001 ],
    [ 0.23986188,  0.28200601,  0.26762535,  0.21050676]]]

    msgsR=[[[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.29002292,  0.2989277 ,  0.24252004,  0.16852933],
    [ 0.26310106,  0.29748655,  0.25634478,  0.18306761],
    [ 0.26094932,  0.29628075,  0.25764968,  0.18512025],
    [ 0.2660311 ,  0.29309694,  0.2554316 ,  0.18544036],
    [ 0.2429187 ,  0.28254909,  0.26607192,  0.20846029]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.37767773,  0.28994844,  0.20252485,  0.12984898],
    [ 0.27528358,  0.30773349,  0.24927874,  0.16770419],
    [ 0.28793716,  0.30515927,  0.24350915,  0.16339442],
    [ 0.34116812,  0.28865721,  0.22147385,  0.14870082],
    [ 0.24021786,  0.28857552,  0.26886635,  0.20234026]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.35925421,  0.30137138,  0.20782554,  0.13154887],
    [ 0.30184483,  0.30957441,  0.23491658,  0.15366418],
    [ 0.31329719,  0.30255431,  0.23159436,  0.15255414],
    [ 0.28215613,  0.30507813,  0.24650166,  0.16626408],
    [ 0.2502795 ,  0.29243356,  0.26351259,  0.19377434]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.43952647,  0.2797989 ,  0.17408407,  0.10659055],
    [ 0.31874805,  0.30821492,  0.22649261,  0.14654442],
    [ 0.29846323,  0.30843893,  0.23701008,  0.15608777],
    [ 0.2975981 ,  0.30546019,  0.23828072,  0.15866099],
    [ 0.2534437 ,  0.29230677,  0.26173494,  0.19251459]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.38851916,  0.28944442,  0.19638156,  0.12565486],
    [ 0.41445041,  0.28153079,  0.1859509 ,  0.1180679 ],
    [ 0.30207854,  0.30005682,  0.23625284,  0.1616118 ],
    [ 0.40847836,  0.27925759,  0.18936394,  0.1229001 ],
    [ 0.23986188,  0.28200601,  0.26762535,  0.21050676]]]

    msgsU=[[[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.29002292,  0.2989277 ,  0.24252004,  0.16852933],
    [ 0.26310106,  0.29748655,  0.25634478,  0.18306761],
    [ 0.26094932,  0.29628075,  0.25764968,  0.18512025],
    [ 0.2660311 ,  0.29309694,  0.2554316 ,  0.18544036],
    [ 0.2429187 ,  0.28254909,  0.26607192,  0.20846029]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.37767773,  0.28994844,  0.20252485,  0.12984898],
    [ 0.27528358,  0.30773349,  0.24927874,  0.16770419],
    [ 0.28793716,  0.30515927,  0.24350915,  0.16339442],
    [ 0.34116812,  0.28865721,  0.22147385,  0.14870082],
    [ 0.24021786,  0.28857552,  0.26886635,  0.20234026]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.35925421,  0.30137138,  0.20782554,  0.13154887],
    [ 0.30184483,  0.30957441,  0.23491658,  0.15366418],
    [ 0.31329719,  0.30255431,  0.23159436,  0.15255414],
    [ 0.28215613,  0.30507813,  0.24650166,  0.16626408],
    [ 0.2502795 ,  0.29243356,  0.26351259,  0.19377434]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.43952647,  0.2797989 ,  0.17408407,  0.10659055],
    [ 0.31874805,  0.30821492,  0.22649261,  0.14654442],
    [ 0.29846323,  0.30843893,  0.23701008,  0.15608777],
    [ 0.2975981 ,  0.30546019,  0.23828072,  0.15866099],
    [ 0.2534437 ,  0.29230677,  0.26173494,  0.19251459]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.38851916,  0.28944442,  0.19638156,  0.12565486],
    [ 0.41445041,  0.28153079,  0.1859509 ,  0.1180679 ],
    [ 0.30207854,  0.30005682,  0.23625284,  0.1616118 ],
    [ 0.40847836,  0.27925759,  0.18936394,  0.1229001 ],
    [ 0.23986188,  0.28200601,  0.26762535,  0.21050676]]]

    msgsD=[[[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.29002292,  0.2989277 ,  0.24252004,  0.16852933],
    [ 0.26310106,  0.29748655,  0.25634478,  0.18306761],
    [ 0.26094932,  0.29628075,  0.25764968,  0.18512025],
    [ 0.2660311 ,  0.29309694,  0.2554316 ,  0.18544036],
    [ 0.2429187 ,  0.28254909,  0.26607192,  0.20846029]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.37767773,  0.28994844,  0.20252485,  0.12984898],
    [ 0.27528358,  0.30773349,  0.24927874,  0.16770419],
    [ 0.28793716,  0.30515927,  0.24350915,  0.16339442],
    [ 0.34116812,  0.28865721,  0.22147385,  0.14870082],
    [ 0.24021786,  0.28857552,  0.26886635,  0.20234026]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.35925421,  0.30137138,  0.20782554,  0.13154887],
    [ 0.30184483,  0.30957441,  0.23491658,  0.15366418],
    [ 0.31329719,  0.30255431,  0.23159436,  0.15255414],
    [ 0.28215613,  0.30507813,  0.24650166,  0.16626408],
    [ 0.2502795 ,  0.29243356,  0.26351259,  0.19377434]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.43952647,  0.2797989 ,  0.17408407,  0.10659055],
    [ 0.31874805,  0.30821492,  0.22649261,  0.14654442],
    [ 0.29846323,  0.30843893,  0.23701008,  0.15608777],
    [ 0.2975981 ,  0.30546019,  0.23828072,  0.15866099],
    [ 0.2534437 ,  0.29230677,  0.26173494,  0.19251459]],
    [[ 1.        ,  1.        ,  1.        ,  1.        ],
    [ 0.38851916,  0.28944442,  0.19638156,  0.12565486],
    [ 0.41445041,  0.28153079,  0.1859509 ,  0.1180679 ],
    [ 0.30207854,  0.30005682,  0.23625284,  0.1616118 ],
    [ 0.40847836,  0.27925759,  0.18936394,  0.1229001 ],
    [ 0.23986188,  0.28200601,  0.26762535,  0.21050676]]]

    msgsL, msgsR, msgsU, msgsD=do_sweepsSS2(unPots, binPots, msgsL, msgsR, msgsU, msgsD, nsweeps, h2, w2, nd, rad, f0)
